import cv2
import json
from collections import defaultdict
from vehicle_flow_counting import CLASS_NAMES

class FlowCounter:
    def __init__(self, frame_shape, class_names=None, edge_margin=100, show_edge=False):
        self.height, self.width = frame_shape
        self.edge_margin = edge_margin
        self.show_edge = show_edge
        self.history = {}       # track_id -> last (cx, cy)

        # 分类别统计 使用 dict[class_id]['enter'/'exit'][edge] 统计
        self.counts = defaultdict(lambda: {
            'enter': defaultdict(int),
            'exit': defaultdict(int)
        })
        # 总统计
        self.total_counts = {
            'enter': defaultdict(int),
            'exit': defaultdict(int)
        }

        self.class_names = class_names if class_names else {}

    # 消失统计法
    # def get_edge(self, cx, cy):
    #     """根据中心点靠近哪条边返回边名"""
    #     if cy < self.edge_margin: return 'top'
    #     elif cy > self.height - self.edge_margin: return 'bottom'
    #     elif cx < self.edge_margin: return 'left'
    #     elif cx > self.width - self.edge_margin: return 'right'
    #     else: return None

    # def update(self, current_tracks):
    #     """
    #     current_tracks: list of tuples (track_id, (cx,cy))
    #     """
    #     current_ids = set()
    #     for track_id, (cx, cy) in current_tracks:
    #         current_ids.add(track_id)
    #         # 新出现 → enter
    #         if track_id not in self.history:
    #             edge = self.get_edge(cx, cy)
    #             if edge:
    #                 self.count_enter[edge] += 1
    #                 print(f"[Enter] ID={track_id} from {edge}")
    #         self.history[track_id] = (cx, cy)

    #     # track消失 → exit
    #     disappeared_ids = set(self.history.keys()) - current_ids
    #     for track_id in disappeared_ids:
    #         cx, cy = self.history[track_id]
    #         edge = self.get_edge(cx, cy)
    #         if edge:
    #             self.count_exit[edge] += 1
    #             print(f"[Exit ] ID={track_id} from {edge}")
    #         del self.history[track_id]

    # def draw_counts(self, frame):
    #     txt = (f"Enter T:{self.count_enter['top']} B:{self.count_enter['bottom']} "
    #            f"L:{self.count_enter['left']} R:{self.count_enter['right']} | "
    #            f"Exit  T:{self.count_exit['top']} B:{self.count_exit['bottom']} "
    #            f"L:{self.count_exit['left']} R:{self.count_exit['right']}")
    #     cv2.putText(frame, txt, (20,30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0,0,255), 2)


    def crossed_edge(self, prev_pos, cur_pos):
        px, py = prev_pos
        cx, cy = cur_pos
        w, h = self.width, self.height
        m = self.edge_margin
        crossed = []

        # top
        if py < m and cy >= m: crossed.append(('top','enter'))
        elif py >= m and cy < m: crossed.append(('top','exit'))

        # bottom
        if py > h - m and cy <= h - m: crossed.append(('bottom','enter'))
        elif py <= h - m and cy > h - m: crossed.append(('bottom','exit'))

        # left
        if px < m and cx >= m: crossed.append(('left','enter'))
        elif px >= m and cx < m: crossed.append(('left','exit'))

        # right
        if px > w - m and cx <= w - m: crossed.append(('right','enter'))
        elif px <= w - m and cx > w - m: crossed.append(('right','exit'))

        return crossed

    def update(self, current_tracks):
        """
        输入: current_tracks = [(track_id, (cx, cy)), ...]
        """
        for track_id, cls_id, cur_center in current_tracks:
            prev_center = self.history.get(track_id)
            if prev_center is not None:
                crossed = self.crossed_edge(prev_center, cur_center)
                for edge, action in crossed:
                    # 分类别
                    self.counts[cls_id][action][edge] += 1
                    # 总统计
                    self.total_counts[action][edge] += 1

                    cname = self.class_names.get(cls_id, str(cls_id))
                    # print(f"[{action.upper()}] ID={track_id} {cname} via {edge}")

            self.history[track_id] = cur_center

    def draw_target(self, frame, current_tracks):
        # 绘制所有目标的跟踪ID、类别和中心点
        for track_id, cls_id, (cx, cy) in current_tracks:
            color = COLORS.get(cls_id, (255, 255, 255))  # 默认白色
            cname = CLASS_NAMES.get(cls_id, str(cls_id))
            cv2.circle(frame, (cx, cy), 5, color, -1)
            cv2.putText(frame, f"ID:{track_id} {cname}", (cx, cy-10),
                cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)

    def draw_counts(self, frame):
        y0 = 60
        h, w = frame.shape[:2]
        # 基于图像宽度或高度动态调整字体缩放与粗细
        font_scale = max(0.5, min(1.6, w / 1000))   # 宽=1000时约1.0
        thickness  = max(1, int(w / 1000))         # 宽=1000时约2

        # 总统计
        txt_total = (f"total | "
                     f"Enter T:{self.total_counts['enter']['top']} "
                     f"B:{self.total_counts['enter']['bottom']} "
                     f"L:{self.total_counts['enter']['left']} "
                     f"R:{self.total_counts['enter']['right']} | "
                     f"Exit T:{self.total_counts['exit']['top']} "
                     f"B:{self.total_counts['exit']['bottom']} "
                     f"L:{self.total_counts['exit']['left']} "
                     f"R:{self.total_counts['exit']['right']}")
        cv2.putText(frame, txt_total, (20, y0), cv2.FONT_HERSHEY_SIMPLEX,
                    font_scale, COLORS.get(18, (255, 255, 255)), thickness)
        y0 += 60

        # 分类别
        for cls_id, stats in self.counts.items():
            cname = self.class_names.get(cls_id, str(cls_id))
            txt = (f"{cname} | "
                   f"Enter T:{stats['enter']['top']} B:{stats['enter']['bottom']} "
                   f"L:{stats['enter']['left']} R:{stats['enter']['right']} | "
                   f"Exit T:{stats['exit']['top']} B:{stats['exit']['bottom']} "
                   f"L:{stats['exit']['left']} R:{stats['exit']['right']}")
            cv2.putText(frame, txt, (20, y0), cv2.FONT_HERSHEY_SIMPLEX,
                font_scale, COLORS.get(cls_id, (255, 255, 255)), thickness)
            y0 += 60

        # 画辅助边界线
        if self.show_edge:
            m = self.edge_margin
            h, w = self.height, self.width
            color = COLORS.get(17, (255, 255, 255))
            cv2.line(frame, (0, m), (w, m), color, 1)          # top
            cv2.line(frame, (0, h - m), (w, h - m), color, 1)  # bottom
            cv2.line(frame, (m, 0), (m, h), color, 1)          # left
            cv2.line(frame, (w - m, 0), (w - m, h), color, 1)  # right

    def get_stats_json(self):
        """返回统计 JSON，total 和各类别在同一层级"""
        data = {"total": {k: dict(v) for k, v in self.total_counts.items()}}  # defaultdict -> dict
        for cls_id, stats in self.counts.items():
            cname = self.class_names.get(cls_id, str(cls_id))
            data[cname] = {k: dict(v) for k, v in stats.items()}
        stats_json = json.dumps(data, indent=2, ensure_ascii=False)
        # print(stats_json)
        return stats_json

    # def draw_and_get_json(self, frame):
    #     """一次性绘制到帧并返回 JSON"""
    #     frame = self.draw_counts(frame)
    #     stats_json = self.get_stats_json()
    #     return frame, stats_json


COLORS = {
    0: (0, 128, 255),       # 橙色 - pedestrian
    1: (255, 0, 128),       # 紫色 - people
    2: (255, 255, 0),       # 青色 - bicycle
    3: (0, 255, 0),         # 绿色 - car
    4: (0, 255, 255),       # 黄色 - van
    5: (255, 0, 255),       # 紫红 - truck
    6: (128, 0, 255),       # 玫红 - tricycle
    7: (0, 255, 128),       # 黄绿 - awning-tricycle
    8: (0, 0, 255),         # 红色 - bus
    9: (255, 0, 0),         # 蓝色 - motor

    10: (128, 255, 0),      # 淡青
    11: (255, 128, 0),      # 淡蓝
    12: (128, 128, 255),    # 粉橙
    13: (128, 255, 128),    # 青绿
    14: (255, 128, 128),    # 淡紫
    15: (128, 255, 255),    # 淡黄
    16: (255, 128, 255),    # 粉色
    17: (255, 255, 128),    # 青蓝
    18: (255, 255, 255),    # 白色
}

